#include "ASMInstruction.hpp"
#include "CodeGen.hpp"
#include "Module.hpp"

#include <iostream>
#include <memory>
#include <unordered_map>

void translate_main(CodeGen *codegen); // 将 main 函数翻译为汇编代码

int main() {
    auto *module = new Module();
    auto *codegen = new CodeGen(module);

    // 告诉汇编器将汇编放到代码段
    codegen->append_inst(".text", ASMInstruction::Atrribute);

    translate_main(codegen);

    std::cout << codegen->print();
    delete codegen;
    delete module;
    return 0;
}

// TODO: 翻译你优化后的IR为RISC-V目标代码
void translate_main(CodeGen *codegen) {
    std::unordered_map<std::string, int> offset_map;

    /* 声明外部函数 */
    codegen->append_inst(".extern getint", ASMInstruction::Atrribute);
    codegen->append_inst(".extern putint", ASMInstruction::Atrribute);

    /* 声明 main 函数 */ 
    codegen->append_inst(".align 2", ASMInstruction::Atrribute);
    codegen->append_inst(".globl main", ASMInstruction::Atrribute);
    codegen->append_inst(".type main, @function", ASMInstruction::Atrribute);

    /* main 函数开始 */
    codegen->append_inst("main", ASMInstruction::Label);

    /* main 函数的 Prologue (序言) */
    codegen->append_inst("addi", {"sp", "sp", "-64"});
    codegen->append_inst("sd", {"ra", "sp", "56"});
    codegen->append_inst("sd", {"s0", "sp", "48"});
    codegen->append_inst("addi", {"s0", "sp", "64"});

    /* main 函数的 label_entry */
    codegen->append_inst(".main_label_entry", ASMInstruction::Label);
    codegen->append_inst("# %op2 = call i32 @getint()", ASMInstruction::Comment);
    codegen->append_inst("call", {"getint"});
    codegen->append_inst("mv", {"s1", "a0"});  // 保存getint的结果到s1 (c的值)
    codegen->append_inst("# br label %label3", ASMInstruction::Comment);
    codegen->append_inst("li", {"t0", "1"});   // 初始化phi节点的值为1 (a的初始值)
    codegen->append_inst("j", {".main_label3"});

    /* label3: 主循环 */
    codegen->append_inst(".main_label3", ASMInstruction::Label);
    codegen->append_inst("# %op4 = phi i32 [ 1, %label_entry ], [ %op_new, %label20 ]", ASMInstruction::Comment);
    codegen->append_inst("# %op5 = icmp slt i32 %op4, 999999999", ASMInstruction::Comment);
    codegen->append_inst("li", {"t1", "999999999"});
    codegen->append_inst("slt", {"t2", "t0", "t1"});
    codegen->append_inst("# br i1 %op5, label %label6, label %label11", ASMInstruction::Comment);
    codegen->append_inst("addi", {"t3", "zero", "0"});
    codegen->append_inst("blt", {"t3", "t2", ".main_label6"});
    codegen->append_inst("j", {".main_label11"});

    /* label6: 检查除法条件 */
    codegen->append_inst(".main_label6", ASMInstruction::Label);
    codegen->append_inst("# %op9 = srem i32 %op4, %op2", ASMInstruction::Comment);
    codegen->append_inst("rem", {"t1", "t0", "s1"});
    codegen->append_inst("# %op10 = icmp eq i32 %op9, 0", ASMInstruction::Comment);
    codegen->append_inst("addi", {"t2", "zero", "0"});
    codegen->append_inst("# br i1 %op10, label %label14, label %label17", ASMInstruction::Comment);
    codegen->append_inst("beq", {"t1", "t2", ".main_label14"});
    codegen->append_inst("j", {".main_label17"});

    /* label11: 循环结束 */
    codegen->append_inst(".main_label11", ASMInstruction::Label);
    codegen->append_inst("# call void @putint(i32 %op2)", ASMInstruction::Comment);
    codegen->append_inst("mv", {"a0", "s1"});
    codegen->append_inst("call", {"putint"});
    codegen->append_inst("# ret i32 %op4", ASMInstruction::Comment);
    codegen->append_inst("mv", {"a0", "t0"});
    codegen->append_inst("j", {".main_exit"});

    /* label14: a%c==0的情况 */
    codegen->append_inst(".main_label14", ASMInstruction::Label);
    codegen->append_inst("# %op16 = add i32 %op4, 3", ASMInstruction::Comment);
    codegen->append_inst("addi", {"t0", "t0", "3"});
    codegen->append_inst("# br label %label20", ASMInstruction::Comment);
    codegen->append_inst("j", {".main_label20"});

    /* label17: a%c!=0的情况 */
    codegen->append_inst(".main_label17", ASMInstruction::Label);
    codegen->append_inst("# %op19 = add i32 %op4, 2", ASMInstruction::Comment);
    codegen->append_inst("addi", {"t0", "t0", "2"});
    codegen->append_inst("# br label %label20", ASMInstruction::Comment);
    codegen->append_inst("j", {".main_label20"});

    /* label20: 返回循环开始 */
    codegen->append_inst(".main_label20", ASMInstruction::Label);
    codegen->append_inst("# %op_new = phi i32 [ %op16, %label14 ], [ %op19, %label17 ]", ASMInstruction::Comment);
    codegen->append_inst("# br label %label3", ASMInstruction::Comment);
    codegen->append_inst("j", {".main_label3"});

    /* main 函数的 Epilogue (收尾) */
    codegen->append_inst(".main_exit", ASMInstruction::Label);
    codegen->append_inst("ld", {"ra", "sp", "56"});
    codegen->append_inst("ld", {"s0", "sp", "48"});
    codegen->append_inst("addi", {"sp", "sp", "64"});
    codegen->append_inst("ret");
}
